简单的 attention 计算模型

如果我们了解attention的基础结构，那么可以了解attention内的基础计算单元的计算过程，由于multihead之间不做交叉运算, 因此简单的看一个head之内的运算, 即QKV矩阵相乘的公式为 :

$$\text{Attention}(Q, K, V) = \text{softmax}!\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

即矩阵 $Q$ 与键矩阵 $K$ 的转置做点积并除以缩放因子 $\sqrt{d_k}$ 后，经 softmax 归一化得到注意力权重，再与值矩阵 $V$ 相乘，即得到对各位置信息加权聚合后的输出。其中缩放因子是方便计算的标量, 因此先简化来看矩阵乘法的视角: 这里假设Q,K,V 的维度都是 (N,d). K,V通常是相同维度,但是可能与 Q维度不同, 具体情况后面再列举.

1. 计算注意力权重矩阵 $P$ $$ P = \operatorname{softmax}!\left(s \cdot Q K^{T},\ \dim=-1\right), \quad s = d^{-\frac{1}{2}} $$ 其中 $QK^{T}$ 形状为 $(N, d) \times (d, N) = (N, N)$，沿最后一维做 softmax 归一化后得到 $P \in \mathbb{R}^{N \times N}$。
2. 计算输出矩阵 $O$ $$ O = P V $$ 其中 $P$ 形状为 $(N, N)$，$V$ 形状为 $(N, d)$，相乘后得到输出 $O \in \mathbb{R}^{N \times d}$。

$$ \boxed{,O = \operatorname{softmax}!\left(\frac{QK^{T}}{\sqrt{d}}\right) V,,\qquad Q,K,V \in \mathbb{R}^{N \times d},\ \ O \in \mathbb{R}^{N \times d},} $$ 值得一提的是, multi-head中的head维度一般不大且固定, 如llama 2 的LLM选择的是128或者56, 方便放入一个 n卡的 block中.

flashatten 的计算有什么改进